package com.xiashitech.agent.instrumentation.spring.security.inst;

import com.xiashitech.agent.cache.status.CacheStatus;
import com.xiashitech.agent.config.StartConfig;
import com.xiashitech.agent.dto.UserRequest;
import com.xiashitech.agent.instrumentation.bridge.Java8Bridge;
import com.xiashitech.agent.instrumentation.utils.BaggageUtil;
import com.xiashitech.agent.utils.DurationPrinter;
import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
import io.opentelemetry.javaagent.shaded.io.opentelemetry.context.Context;
import io.opentelemetry.javaagent.shaded.io.opentelemetry.context.Scope;
import net.bytebuddy.asm.Advice;
import net.bytebuddy.description.type.TypeDescription;
import net.bytebuddy.matcher.ElementMatcher;
import org.springframework.security.core.Authentication;

import java.util.Date;

import static com.xiashitech.agent.instrumentation.spring.security.XSSpringSecuritySingletons.instrumenter;
import static io.opentelemetry.javaagent.bootstrap.Java8BytecodeBridge.currentContext;
import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.implementsInterface;
import static net.bytebuddy.matcher.ElementMatchers.*;

public class XSSpringSecurityContextInstrumentation implements TypeInstrumentation {
    @Override
    public ElementMatcher<TypeDescription> typeMatcher() {
        return implementsInterface(named("org.springframework.security.core.context.SecurityContext"));
    }

    @Override
    public void transform(TypeTransformer transformer) {
        transformer.applyAdviceToMethod(
                isMethod()
                        .and(isPublic())
                        .and(named("getAuthentication"))
                        .and(takesArguments(0))
                        .and(returns(named("org.springframework.security.core.Authentication"))),
                XSSpringSecurityContextInstrumentation.class.getName() + "$AuthenticateAdvice");
    }
    public static class AuthenticateAdvice {

        @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
        public static void onExit(
                @Advice.Return Authentication authentication,
                @Advice.Local("otelContext") Context context,
                @Advice.Local("otelScope") Scope scope) {
            DurationPrinter durationPrinter = new DurationPrinter("XSSpringSecurityContextInstrumentation");
            if(authentication == null) {
                return;
            }
            if (CacheStatus.cannotRun()) {
                return;
            }
            String uid = Java8Bridge.currentForBaggage().getEntryValue("xs.userId");
            if (uid != null && !"".equals(uid) && !"anonymousUser".equalsIgnoreCase(uid) && !"pcsa".equalsIgnoreCase(uid)){
                return;
            }
            try {
                UserRequest userRequest = UserRequest.builder()
                        .requestDate(String.valueOf(new Date().getTime()))
                        .loginRequest(Boolean.TRUE)
                        .build();
                Context parentContext = currentContext();
                if (instrumenter().shouldStart(parentContext,userRequest)) {
                    parentContext = BaggageUtil.syncBaggage().build().storeInContext(parentContext);
                    context = instrumenter().start(parentContext, userRequest);
                    context = BaggageUtil.syncBaggage().build().storeInContext(context);
                    scope = context.makeCurrent();

                    // 生成 span
                    scope.close();
                    instrumenter().end(context, userRequest, authentication,null);
                }
            } catch (Throwable e) {
                StartConfig.throwableStackPrinter(e);
            }
            durationPrinter.printElapseDuration("AuthenticateAdvice e");
        }
    }
}
